import os
import json
from pathlib import Path
from pycocotools.coco import COCO


def convert_coco_to_yolo(coco_json_path, images_dir, output_dir):
    """
    Convert COCO annotation to YOLO format.

    Args:
        coco_json_path (str): Path to COCO JSON annotation file.
        images_dir (str): Path to directory containing images.
        output_dir (str): Path to save YOLO formatted .txt files.
    """
    coco = COCO(coco_json_path)
    cat_ids = coco.getCatIds()
    categories = coco.loadCats(cat_ids)
    category_mapping = {cat['id']: i for i, cat in enumerate(categories)}

    Path(output_dir).mkdir(parents=True, exist_ok=True)

    img_ids = coco.getImgIds()
    for img_id in img_ids:
        img_info = coco.loadImgs(img_id)[0]
        img_filename = img_info['file_name']
        img_width = img_info['width']
        img_height = img_info['height']

        ann_ids = coco.getAnnIds(imgIds=img_id)
        anns = coco.loadAnns(ann_ids)

        yolo_annotations = []
        for ann in anns:
            category_id = category_mapping[ann['category_id']]
            bbox = ann['bbox']
            x, y, w, h = bbox
            x_center = (x + w / 2) / img_width
            y_center = (y + h / 2) / img_height
            width = w / img_width
            height = h / img_height
            yolo_annotations.append(f"{category_id} {x_center} {y_center} {width} {height}")

        txt_filename = os.path.splitext(img_filename)[0] + ".txt"
        with open(os.path.join(output_dir, txt_filename), "w") as f:
            f.write("\n".join(yolo_annotations))


# 示例调用
convert_coco_to_yolo(
    coco_json_path="_annotations.coco.json",
    images_dir="train2017",
    output_dir="../data/worker_safe/labels/train"
)
